import Image
import numpy
import os
import operator

PHOTO_PREFIX = "s"

class TrainingSet:

	def __init__(self, directory):
		self.labels = []
		images = []
		for facedir in os.listdir(directory):
			if facedir.startswith(PHOTO_PREFIX):
				facepath = os.path.join(directory, facedir)
				for imgfile in os.listdir(facepath):
					if not imgfile.startswith('.'):
						self.labels.append(facedir)
						images.append(Image.open(os.path.join(facepath, imgfile)))
		self.matrix = numpy.empty((len(images), reduce(operator.mul, images[0].size)))
		for i in range(len(images)):
			self.matrix[i] = numpy.array(images[i].getdata())

class TestSet:

	def __init__(self, directory):
		self.labels = []
		images = []
		for imgfile in os.listdir(directory):
			if (imgfile.startswith(PHOTO_PREFIX)):
				self.labels.append(os.path.splitext(imgfile)[0])
				images.append(Image.open(os.path.join(directory, imgfile)))
		self.matrix = numpy.empty((len(images), reduce(operator.mul, images[0].size)))
		for i in range(len(images)):
			self.matrix[i] = numpy.array(images[i].getdata())
